# %%
import numpy as np
import scipy.stats as stats
from numpy import sqrt, exp
import scipy.optimize as opt
import matplotlib.pyplot as plt

# import numba
import math
import warnings

warnings.filterwarnings("ignore")  # 设置警告过滤器，忽略所有警告
plt.rcParams["font.sans-serif"] = ["SimHei"]  # 用来正常显示中文标签
plt.rcParams["axes.unicode_minus"] = False  # 正常显示负号
ITERATION_MAX_ERROR = 0.00001  # 牛顿法迭代的精度

# %% [markdown]
# 1.BAW期权定价模型


# %%
# BSM期权定价函数
def BSM(CP, S, X, sigma, T, r, b):
    """
    Parameters
    ----------
    CP：看涨或看跌"C"or"P"
    S : 标的价格.
    X : 行权价格.
    sigma :波动率.
    T : 年化到期时间. 即，到期时间（以天数为单位）/252
    r : 收益率.
    b : 持有成本，当b = r 时，为标准的无股利模型，b=0时，为期货期权，b为r-q时，为支付股利模型，b为r-rf时为外汇期权.
    Returns
    -------
    返回欧式期权的估值
    """
    d1 = (np.log(S / X) + (b + sigma**2 / 2) * T) / (sigma * np.sqrt(T))
    d2 = d1 - sigma * np.sqrt(T)
    if CP == "C":
        value = S * np.exp((b - r) * T) * stats.norm.cdf(d1) - X * np.exp(
            -r * T
        ) * stats.norm.cdf(
            d2
        )  # stats.norm.cdf() 是 SciPy 库中 stats 模块的一个函数，用于计算标准正态分布的累积分布函数 (CDF)。CDF 返回的是一个给定值小于或等于某个特定值的概率。
    else:
        value = X * np.exp(-r * T) * stats.norm.cdf(-d2) - S * np.exp(
            (b - r) * T
        ) * stats.norm.cdf(-d1)
    return value


# %%
# 牛顿迭代法求解S的函数
def find_Sx(CP, X, sigma, T, r, b):  # 手动写的标准的牛顿迭代法
    M = 2 * r / sigma**2
    N = 2 * b / sigma**2
    K = 1 - exp(-r * T)
    q1 = (-(N - 1) - sqrt((N - 1) ** 2 + 4 * M / K)) / 2
    q2 = (-(N - 1) + sqrt((N - 1) ** 2 + 4 * M / K)) / 2
    if CP == "C":
        # 初始化Si的值
        S_infinite = X / (
            1 - 2 * (-(N - 1) + sqrt((N - 1) ** 2 + 4 * M)) ** -1
        )  # 到期时间为无穷时的价格
        h2 = -(b * T + 2 * sigma * sqrt(T)) * X / (S_infinite - X)
        Si = X + (S_infinite - X) * (1 - exp(h2))  # 计算种子值
        # print(f"Si的种子值为{Si}")
        LHS = Si - X
        d1 = (np.log(Si / X) + (b + sigma**2 / 2) * T) / (sigma * sqrt(T))
        RHS = (
            BSM("C", Si, X, sigma, T, r, b)
            + (1 - exp((b - r) * T) * stats.norm.cdf(d1)) * Si / q2
        )
        bi = (
            exp((b - r) * T) * stats.norm.cdf(d1) * (1 - 1 / q2)
            + (1 - (exp((b - r) * T) * stats.norm.pdf(d1)) / sigma / sqrt(T)) / q2
        )  # bi为迭代使用的初始斜率（RHS的初始导数）
        while (
            np.abs((LHS - RHS) / X) > ITERATION_MAX_ERROR
        ):  # 通过迭代求解满足精度要求的Si
            Si = (X + RHS - bi * Si) / (1 - bi)  # 牛顿迭代公式
            # print(f"Si的值迭代为{Si}")
            LHS = Si - X
            d1 = (np.log(Si / X) + (b + sigma**2 / 2) * T) / (sigma * sqrt(T))
            RHS = (
                BSM("C", Si, X, sigma, T, r, b)
                + (1 - exp((b - r) * T) * stats.norm.cdf(d1)) * Si / q2
            )
            bi = (
                exp((b - r) * T) * stats.norm.cdf(d1) * (1 - 1 / q2)
                + (1 - (exp((b - r) * T) * stats.norm.pdf(d1)) / sigma / sqrt(T)) / q2
            )
        return Si
    else:  # 看跌期权与看涨期权类似
        S_infinite = X / (1 - 2 * (-(N - 1) - sqrt((N - 1) ** 2 + 4 * M)) ** -1)
        h1 = -(b * T - 2 * sigma * sqrt(T)) * X / (X - S_infinite)
        Si = S_infinite + (X - S_infinite) * exp(h1)  # 计算种子值
        # print(f"Si的种子值为{Si}")
        LHS = X - Si
        d1 = (np.log(Si / X) + (b + sigma**2 / 2) * T) / (sigma * sqrt(T))
        RHS = (
            BSM("P", Si, X, sigma, T, r, b)
            - (1 - exp((b - r) * T) * stats.norm.cdf(-d1)) * Si / q1
        )
        bi = (
            -exp((b - r) * T) * stats.norm.cdf(-d1) * (1 - 1 / q1)
            - (1 + (exp((b - r) * T) * stats.norm.pdf(-d1)) / sigma / sqrt(T)) / q1
        )
        while np.abs((LHS - RHS) / X) > ITERATION_MAX_ERROR:
            Si = (X - RHS + bi * Si) / (1 + bi)
            # print(f"Si的值迭代为{Si}")
            LHS = X - Si
            d1 = (np.log(Si / X) + (b + sigma**2 / 2) * T) / (sigma * sqrt(T))
            RHS = (
                BSM("P", Si, X, sigma, T, r, b)
                - (1 - exp((b - r) * T) * stats.norm.cdf(-d1)) * Si / q1
            )
            bi = (
                -exp((b - r) * T) * stats.norm.cdf(-d1) * (1 - 1 / q1)
                - (1 + (exp((b - r) * T) * stats.norm.pdf(-d1)) / sigma / sqrt(T)) / q1
            )
        return Si


# %%
# 使用了SciPy库中的优化函数 opt.fmin 来迭代求解期权的隐含价格Si
def find_Sx_func(CP, S, X, sigma, T, r, b):  # 用来计算目标函数的值
    M = 2 * r / sigma**2
    N = 2 * b / sigma**2
    K = 1 - exp(-r * T)
    q1 = (-(N - 1) - sqrt((N - 1) ** 2 + 4 * M / X)) / 2
    q2 = (-(N - 1) + sqrt((N - 1) ** 2 + 4 * M / X)) / 2
    if CP == "C":
        LHS = S - X
        RHS = (
            BSM("C", S, X, sigma, T, r, b)
            + (
                1
                - exp((b - r) * T)
                * stats.norm.cdf(
                    (np.log(S / X) + (b + sigma**2 / 2) * T) / (sigma * sqrt(T))
                )
            )
            * S
            / q2
        )
        y = (RHS - LHS) ** 2  # 目标函数y=(RHS-LHS)^2
    else:
        LHS = X - S
        RHS = (
            BSM("P", S, X, sigma, T, r, b)
            - (
                1
                - exp((b - r) * T)
                * stats.norm.cdf(
                    -((np.log(S / X) + (b + sigma**2 / 2) * T) / (sigma * sqrt(T)))
                )
            )
            * S
            / q1
        )
        y = (RHS - LHS) ** 2
    return y


def find_Sx_opt(CP, S, X, sigma, T, r, b):  # 优化目标函数，找到使得目标函数最小化的Sii
    start = S  # 随便给一个S的初始值，或者其他值都行
    func = lambda S: find_Sx_func(
        CP, S, X, sigma, T, r, b
    )  # func 是一个匿名函数，用于包装目标函数 find_Sx_func
    Si = opt.fmin(
        func, start, xtol=1e-8, ftol=1e-8
    )  # 直接做掉包侠,xtol=1e-8 和 ftol=1e-8 分别设置了参数和目标函数值的容忍度为 1e-8
    return Si


# %%
# 定义BAW期权定价函数
def BAW(CP, S, X, sigma, T, r, b, opt_method="newton"):
    if b > r:  # b>r时，美式期权价值和欧式期权相同
        value = BSM(CP, S, X, sigma, T, r, b)

    else:
        M = 2 * r / sigma**2
        N = 2 * b / sigma**2
        K = 1 - exp(-r * T)
        if opt_method == "newton":  # 若为牛顿法就用第一种迭代
            Si = find_Sx(CP, X, sigma, T, r, b)
        else:  # 若不为牛顿法，其他方法这里就是scipy的优化方法
            Si = find_Sx_opt(CP, S, X, sigma, T, r, b)
        d1 = (np.log(Si / X) + (b + sigma**2 / 2) * T) / (sigma * sqrt(T))
        if CP == "C":  # 看涨期权
            q2 = (-(N - 1) + sqrt((N - 1) ** 2 + 4 * M / K)) / 2
            A2 = Si / q2 * (1 - exp((b - r) * T) * stats.norm.cdf(d1))
            if S < Si:
                value = BSM(CP, S, X, sigma, T, r, b) + A2 * (S / Si) ** q2
            else:
                value = S - X

        else:  # 看跌期权
            q1 = (-(N - 1) - sqrt((N - 1) ** 2 + 4 * M / K)) / 2
            A1 = -Si / q1 * (1 - exp((b - r) * T) * stats.norm.cdf(-d1))
            if S > Si:
                value = BSM(CP, S, X, sigma, T, r, b) + A1 * (S / Si) ** q1
            else:
                value = X - S
    return value


# %%
BAW(
    CP="P", S=100, X=99, sigma=0.2, T=1, r=0.03, b=0, opt_method="newton"
)  # 使用牛顿迭代法

# %%
BAW(
    CP="P", S=100, X=99, sigma=0.2, T=1, r=0.03, b=0, opt_method="scipy"
)  # 使用优化函数 opt.fmin

# %% [markdown]
# 2.比较不同期权时BAW和二叉树的计算结果


# %%
# 二叉树模型
def simulate_tree_am(CP, m, S0, T, sigma, K, r, b):  # 二叉树模型美式期权
    """
    CP:看涨或看跌
    m：模拟的期数
    S0：期初价格
    T：期限
    sigma：波动率
    K：行权价格
    r:无风险利率
    b:持有成本,当b = r 时，为标准的无股利模型，b=0时，为black76，b为r-q时，为支付股利模型，b为r-rf时为外汇期权
    """
    dt = T / m
    u = math.exp(sigma * math.sqrt(dt))
    d = 1 / u
    S = np.zeros((m + 1, m + 1))
    S[0, 0] = S0
    p = (math.exp(b * dt) - d) / (u - d)
    for i in range(1, m + 1):  # 模拟每个节点的价格
        for a in range(i):
            S[a, i] = S[a, i - 1] * u
            S[a + 1, i] = S[a, i - 1] * d
    Sv = np.zeros_like(S)  # 创建期权价值的矩阵，用到从最后一期倒推期权价值
    if CP == "C":
        S_intrinsic = np.maximum(S - K, 0)
    else:
        S_intrinsic = np.maximum(K - S, 0)
    Sv[:, -1] = S_intrinsic[:, -1]
    for i in range(m - 1, -1, -1):  # 反向倒推每个节点的价值
        for a in range(i + 1):
            Sv[a, i] = max(
                (Sv[a, i + 1] * p + Sv[a + 1, i + 1] * (1 - p)) / np.exp(r * dt),
                S_intrinsic[a, i],
            )
    return Sv[0, 0]


# %%
# 定义模型比较的函数
def model_compare(T_array):
    baw_value = np.zeros_like(T_array)
    crr_value = np.zeros_like(T_array)
    bsm_value = BSM(
        CP="P", S=100, X=99, sigma=0.2, T=T_array, r=0.03, b=0.03
    )  # BSM模型计算的欧式期权价格
    for i in range(len(T_array)):
        baw_value[i] = BAW(
            CP="P",
            S=100,
            X=99,
            sigma=0.2,
            T=T_array[i],
            r=0.03,
            b=0.03,
            opt_method="newton",
        )  # BAW模型计算的美式期权价格
        crr_value[i] = simulate_tree_am(
            CP="P", m=1000, S0=100, T=T_array[i], sigma=0.2, K=99, r=0.03, b=0.03
        )  # CRR 树模型计算的美式期权价格
    plt.plot(T_array, baw_value, label="BAW")
    plt.plot(T_array, crr_value, label="CRR")
    plt.plot(T_array, bsm_value, label="BSM")
    plt.legend()


# %%
model_compare(T_array=np.linspace(0.01, 100, 50))  # 比较100年以内各种模型的定价

# %%
model_compare(T_array=np.linspace(0.01, 1, 20))  # 比较1年以内各种模型的定价
